#!/usr/bin/env python3
# cmp_dir.py
import os
import sys
from collections import defaultdict


def scan(root):
    """
    递归扫描 root 目录，返回 dict:
        key   -> 相对路径（相对于 root）
        value -> 文件大小
    """
    index = {}
    for dirpath, _, files in os.walk(root):
        for f in files:
            abs_path = os.path.join(dirpath, f)
            rel_path = os.path.relpath(abs_path, root)
            try:
                index[rel_path] = os.path.getsize(abs_path)
            except OSError:
                # 文件不可读，记为 None
                index[rel_path] = None
    return index


def compare(dirA, dirB):
    idxA = scan(dirA)
    idxB = scan(dirB)

    # 1. 只在 A 或只在 B 的文件
    only_in_A = set(idxA) - set(idxB)
    only_in_B = set(idxB) - set(idxA)

    # 2. 两边都有的文件，但大小不同
    diff_size = []
    for p in set(idxA) & set(idxB):
        if idxA[p] != idxB[p]:
            diff_size.append((p, idxA[p], idxB[p]))

    return only_in_A, only_in_B, diff_size


def main():
    # if len(sys.argv) != 3:
    #     print("Usage: python cmp_dir.py <dir1> <dir2>")
    #     sys.exit(1)
    #
    # dir1, dir2 = map(os.path.abspath, sys.argv[1:])

    dir1 = r'D:\01.量化\国信iquant\pandas'
    dir2 = r'D:\Future\国信iQuant策略交易平台26\bin.x64\Lib\site-packages\pandas'


    if not (os.path.isdir(dir1) and os.path.isdir(dir2)):
        print("Both arguments must be directories")
        sys.exit(1)

    only1, only2, diff = compare(dir1, dir2)

    if only1:
        print("Only in", dir1)
        for f in sorted(only1):
            print("  ", f)

    if only2:
        print("Only in", dir2)
        for f in sorted(only2):
            print("  ", f)

    if diff:
        print("Different size (file  sizeA  sizeB)")
        for f, s1, s2 in sorted(diff):
            print(f"  {f}  {s1}  {s2}")

    if not any((only1, only2, diff)):
        print("Directories are identical (filename + size) ✓")


if __name__ == "__main__":
    main()
